# TODO: Add comment
# 
# Author: yaping
###############################################################################
library(gplots)
library(colorRamps)
library(RColorBrewer)
##hclust in this package is much faster, and allow more rows...
library(fastcluster)

jet.colors <-colorRampPalette(c("#00007F", "blue", "#007FFF", "cyan","#7FFF7F", "yellow", "#FF7F00", "red", "#7F0000"))
white2red<-colorRampPalette(c("white","red"))
subClusterColors<-c("red","orange","blue","black","cyan","purple","palegreen3","yellow","lightpink2","ivory4")
#mypalette<-colorRampPalette("white","red")
#col=mypalette(100
#image(1:100,1,as.matrix(1:100), col=mypalette(100),xlab="Greens (sequential)",ylab="",xaxt="n",yaxt="n",bty="n"))

generateHeatmapForChipSeqFromHomerOutput<-function(fileName,prefix="heatmap",outputAllInOne=TRUE,sampleLen=101, scale=1000, move_step=20, sample_names,onlyTheseSamples=c(1:length(sample_names)),rowOrder, 
		capUplimit=NULL, capDownLimit=NULL,pngOutput=FALSE,pdfOutput=TRUE, logScale=FALSE,
		averagePlot=FALSE, subClusterNum=1, subClusterOrder=NULL, heatMapCols=white2red(75),
		...){
	heatmap.data<-read.table(fileName,sep="\t",header=T)
	rownames(heatmap.data)<-heatmap.data[,1]
	heatmap.data.order<-heatmap.data[rowOrder,]
	heatmap.data.order<-heatmap.data.order[,2:length(heatmap.data.order[1,])]
	#mat<-array()
	#for(i in 2:length(heatmap.data.order[1,]) ){
	#	mat<-cbind(mat,as.numeric(heatmap.data.order[,i]))
	#}
	#mat<-mat[,2:length(mat[1,])]

	matNorm<-array()
	for(i in onlyTheseSamples){
	#for(i in c(1:1)){
		j=(i-1)*sampleLen+1
		matNorm<-cbind(matNorm,heatmap.data.order[,j:(j+sampleLen-1)])

	}
	matNorm<-matNorm[,2:length(matNorm[1,])]
	if(is.null(capUplimit)){
		capUplimit=quantile(unlist(matNorm),probs=seq(0,1,0.05))[20] #cap top 5%
		
	}
	if(is.null(capDownLimit)){
		capDownLimit=quantile(unlist(matNorm),probs=seq(0,1,0.05))[2] #cap down 5%
		capDownLimit=ifelse(capDownLimit<0,0,capDownLimit)
	}

	matNormCap<-apply(matNorm, 1:2, function(x) ifelse(x<capDownLimit,capDownLimit,ifelse(x>capUplimit,capUplimit,x)))
	#matNormCap<-ifelse(matNorm<capDownLimit,capDownLimit,ifelse(matNorm>capUplimit,capUplimit,matNorm))
	
	if(subClusterNum > 0 & averagePlot){
		axisSeq<-seq(0-scale, scale, by=move_step)
		axisSeqForPlot<-seq(0-scale, scale, by=scale/2)
		pdf(paste("AveragePlot.",subClusterNum,"clusters.",prefix,".",fileName,".chipSeq.pdf",sep=""))
		for(clusterNum in c(1:subClusterNum)){
			cluster_name<-names(subClusterOrder[subClusterOrder==clusterNum])
			valueTmpCluster<-matNormCap[cluster_name,onlyTheseSamples]
			valueGch1<-array()
			for(j in dataSeq){	
				valueGch1<-cbind(valueGch1,mean(colMeans(valueTmpCluster[,(j-move_step/2):(j+move_step/2-1)], na.rm=T), na.rm=T))
			}
			valueGch1<-valueGch1[,2:length(valueGch1[1,])]
			#valueGch1<-colMeans(valueTmpCluster,na.rm=T)			
			
			plot(axisSeq,valueGch1,type="l",axes=FALSE,ylim=c(0,capUplimit),xlab="",ylab="",col=subClusterColors[clusterNum],lty=1,font=2,lwd=3)
			par(new=T)
		}
		axis(1,at=axisSeqForPlot,lty=1,font=2,cex.axis=1.2,cex.lab=1.2,font.lab=2,lwd=2)
		axis(2,at=seq(0,capUplimit,by=capUplimit/5),labels=format(seq(0,capUplimit,by=capUplimit/5),digits=3),lty=1,font=2,cex.axis=1.0,cex.lab=0.8,font.lab=2,lwd=2,las=1)
		title(paste("AveragePlot.",subClusterNum,"clusters.",prefix,"\n",fileName,sep=""), cex.main = 0.6, font.main= 4, col.main= "black",xlab="Distance to center (bp)", ylab="Normalized tag density")
		abline(v=0)
		dev.off()
	}
	


	#j=1
	figureWidth=ifelse(outputAllInOne, 150*length(onlyTheseSamples),480) 
	numFiles=ifelse(outputAllInOne,c(1:1),onlyTheseSamples)

	for(i in numFiles){
		name<-sample_names[i]
		j=(i-1)*sampleLen+1
		tmp<-matNormCap[,j:(j+sampleLen-1)]
		if(outputAllInOne){
			tmp<-matNormCap
		}
		if(pdfOutput){
			fileNamePdf=paste(prefix,".",ifelse(outputAllInOne,"all_samples",name),".pdf",sep="")		
			pdf(fileNamePdf,width = figureWidth, height = ifelse(outputAllInOne,960,480))
		#heatmap.2(ifelse(outputAllInOne,matNormCap, matNormCap[,j:(j+100)]),
			heatmap.3(tmp,
				Rowv=FALSE,
				Colv=FALSE,
				dendrogram= c("none"),
				na.rm=TRUE,
				labRow = "",
				labCol = "",
				#key=TRUE,
				keysize=1,
				trace="none",
				density.info=c("none"),
				margins=c(10, 8),
				na.color=par("bg"),
				margins=c(1, 1),
				col=heatMapCols,
				addAverage=addAverageToHeatmap,
				averageData=content,
				clusterNums=subClusterNum,
				clusterOrder=subClusterOrder,
				colAveragePlot=subClusterColors,
				xAxisSeqForAvePlot=axisSeqForPlot,
				yAxisSeqForAvePlot=seq(0,1,by=0.2),
				axisSeq=axisSeq,
				dataStep=move_step,
				dataSeq=dataSeq,
					...
#col=jet.colors
			)
			dev.off()
		}
		if(pngOutput){
			fileNamePng=paste(prefix,".",ifelse(outputAllInOne,"all_samples",name),".png",sep="")	
			png(fileNamePng,width = figureWidth, height = ifelse(outputAllInOne,960,480))
			heatmap.3(tmp,
				Rowv=FALSE,
				Colv=FALSE,
				dendrogram= c("none"),
				na.rm=TRUE,
				labRow = "",
				labCol = "",
				key=TRUE,
				keysize=1,
				trace="none",
				density.info=c("none"),
				margins=c(10, 8),
				na.color=par("bg"),
				#col=whiteRed.color(10)
				...
			)
			dev.off()
		}
		
	}
	
	returnResult<-list(capUplimit=capUplimit, capDownLimit=capDownLimit) ##breaks is number of col+1
	return(returnResult)
}